#pragma once
#include "wavAudio.h"

  // 估算可用内存的函数
int wavMemory() {
  char *buf;
  int size = 500000; // 可根据需要调整
  while ((buf = (char *)malloc(size -= 1000)) == NULL);
  free(buf);
  return size;
}

// TODO
void wavAudio::readInfo(const char* filepath)
{
  Serial.println("begin to read info...");

  File audio = SD.open(filepath);
  if (!audio) {
    Serial.print("Failed to open WAV file: ");Serial.println(filepath);
    while (1);
  }

  Serial.println("Reading audioFile header ...");
  WAVHeader header;
  audio.read(&header, sizeof(header));
  this->format = header.audioFormat;
  this->channels = header.numChannels;
  this->sample_rate = header.sampleRate;
  this->bits_per_sample = header.bitsPerSample;
  this->data_length = header.subchunk2Size;

  Serial.println("this->samples_num before");
  this->samples_num = this->data_length * 8 / bits_per_sample;

  Serial.print("std::vector<int16_t> before Free memory: ");Serial.println(wavMemory());
  Serial.print("need memory: ");Serial.println(this->samples_num*2);
  // this->data = std::vector<int16_t>(this->samples_num, 0);
  // this->process_data = std::vector<int16_t>(this->samples_num, 0);
  this->data = Malloc(int16_t, this->samples_num);
  this->process_data = Malloc(int16_t, this->samples_num);
  Serial.print("std::vector<int16_t> after Free memory: ");Serial.println(wavMemory());

  Serial.println("audio.read before");
  // audio.read(&this->data, sizeof(this->data));
  // audio.read(this->data.data(), this->data.size() * sizeof(int16_t));
  // audio.read(this->data, this->samples_num * sizeof(int16_t)); 一次性读太多数据会导致只有前面小部分数据正确，后面数据混乱错误

  int k = 0;
  while (audio.available()) {
    audio.read(&this->data[k++], 2);
  }
  audio.close();

  // 打印WAV文件头信息
  Serial.println("WAV File Header Information:");
  Serial.print("Audio Format: "); Serial.println(this->format);
  Serial.print("Number of Channels: "); Serial.println(this->channels);
  Serial.print("Sample Rate: "); Serial.println(this->sample_rate);
  Serial.print("Bits per Sample: "); Serial.println(this->bits_per_sample);
  Serial.print("Data Length (bytes): "); Serial.println(this->data_length);
  
  // 计算并打印样本数
  Serial.print("Number of Samples: "); Serial.println(this->samples_num);


  Serial.println("WAV File Data Information:");
  for (int i = 0; i < 100; i++) {
    Serial.print(this->data[i]);
    if (i % 10 == 1) {
      Serial.println();
    } else {
      Serial.print(" ");
    }
  }

  Serial.println("read info done");
}

bool wavAudio::processdata()
{
  short max_amplitude = 0;
  short min_amplitude = 65535;
  long long mean_amplitude = 0;
  int num = 0;
  for (int i = 0; i < this->samples_num; i++) {
  // for (int i = 0; i < this->data.size(); i++) {
    this->max_data = this->max_data > this->data[i] ? this->max_data : this->data[i];
    max_amplitude = max_amplitude < this->data[i] ? this->data[i] : max_amplitude;
    min_amplitude = min_amplitude > this->data[i] ? this->data[i] : min_amplitude;
    mean_amplitude += abs(this->data[i]);
    num++;
    // Serial.print(this->data[i]);Serial.print(" ");
  }
  Serial.println(" ");
  Serial.print("max_amplitude: ");Serial.println(max_amplitude);
  Serial.print("min_amplitude: ");Serial.println(min_amplitude);
  Serial.print("mean_amplitude: ");Serial.print(mean_amplitude);Serial.print(" / ");Serial.print(num);Serial.print(" = ");Serial.println(mean_amplitude/num);
  Serial.print("diff_amplitude: ");Serial.println(max_amplitude-min_amplitude);
  mean_amplitude = mean_amplitude/num;
  if(!(mean_amplitude > 1000 && min(abs(max_amplitude), abs(min_amplitude)) > 10000)){
    return false;
  }
  
  
  this->_per = int(this->max_data * 0.1);
  Serial.println("max_data done");

  int index = 0; // record the third num(not 0) position
  for (int i = 0; i < this->samples_num; i++) {
  // for (int i = 0; i < this->data.size(); i++) {
    if (abs(this->data[i]) <= this->_per) {
      this->process_data[i] = 0;
    }
    else {
      this->process_data[i] = int(this->data[i] * PEAK / this->max_data);
      index++;
      if (index == 4) {
        this->begin_clock = i;
      }
    }
  }
  Serial.println("process_data done");

  index = 0; // record the end position
  for (int i = this->samples_num - 1; i >= 0; i--) {
  // for (int i = this->process_data.size() - 1; i >= 0; i--) {
    if (this->process_data[i] != 0) {
      index++;
      if (index == 4) {
        this->end_clock = i;
        break;
      }
    }
  }

  return true;
}

void wavAudio::calculatePitch()
{
  int index[100] = { 0 };
  for (int i = this->begin_clock; i < this->end_clock; i++) {
    if (this->process_data[i] != 0) {
      index[this->process_data[i] / 100 - 1]++;
    }
  }
  int n = 0;
  int m = INT8_MIN;
  for (int i = 0; i < 100; i++) {
    if (index[i] > m) {
      n = i;
      m = index[i];
    }
  }
  if (n != 0) {
    this->pitch = n * 100 + 50;
  } else {
    this->pitch = 0;
  }
}

void wavAudio::calculateAverageLoudness()
{
  int num = 0;
  int sum = 0;
  for (int i = this->begin_clock; i < this->end_clock; i++) {
    if (this->process_data[i] != 0) {
      num++;
      sum += this->process_data[i] ^ 2;
    }
  }
  if (num != 0) {
    this->average_loudness = int(sqrt(sum / num));
  }
  else {
    this->average_loudness = 0;
  }
}

void wavAudio::calculateTime()
{
  this->duration = float(this->end_clock - this->begin_clock) / float(this->sample_rate * this->channels);
}

void wavAudio::calculateEnergy()
{
  double energy = 0;
  float d = 1 / float(this->sample_rate * this->channels);
  for (int i = this->begin_clock; i < this->end_clock; i++) {
    if (this->process_data[i] != 0) {
      energy += (this->process_data[i] ^ 2) * d;
    }
  }
  this->energy = energy;
}

void wavAudio::calculateSTFT()
{
  // std::vector<float> x(this->samples_num);
  // std::vector<float> x(this->process_data.size());
  float* x = Malloc(float, this->samples_num);
  // std::transform(this->process_data.begin(), this->process_data.end(), x.begin(),
  //  [](int16_t a) {
  //    return static_cast<float>(a) / 32767.f;
  //  });
  std::transform(this->process_data, this->process_data+this->samples_num, x,
    [](int16_t a) {
      return static_cast<float>(a) / 32767.f;
    });
  int n_fft = 400;
  int n_hop = 160;
  int n_mel = 40;
  int fmin = 80;
  int fmax = 7600;

  Serial.println("transform after");

  //std::vector<std::vector<std::complex<float>>> X = librosa::Feature::stft(x, n_fft, n_hop, "hann", true, "reflect");

  std::vector<float> x1;
  // for (int i = 0, index = 0; i < DATA_NUM && index < x.size(); index++) {
    for (int i = 0, index = 0; i < DATA_NUM && index < this->samples_num; index++) {
    if (x[index] != 0) {
      x1.push_back(x[index]);
      i++;
    }
  }
  std::vector<std::vector<std::vector<float>>> X = simplifiedSTFT(x1, n_fft, n_hop, "hann", true, "reflect");
  Serial.println("simplifiedSTFT after");
  
  // 将实部和虚部分别存储在不同的数组中
  std::vector<float> real_part;
  std::vector<float> imag_part;
  for (std::vector<std::vector<float>> row : X) {
    for (std::vector<float> val : row) {
      real_part.push_back(val[0]);
      imag_part.push_back(val[1]);
    }
  }
  Serial.println("imag_part after");

  // 计算实部的平均值、方差、最大值和最小值
  this->sftf_real_mean = std::accumulate(real_part.begin(), real_part.end(), 0.0f) / real_part.size();
  float sum_of_squares_real = std::inner_product(real_part.begin(), real_part.end(), real_part.begin(), 0.0f);
  this->sftf_real_variance = sum_of_squares_real / real_part.size();
  this->sftf_real_min = *std::min_element(real_part.begin(), real_part.end());
  this->sftf_real_max = *std::max_element(real_part.begin(), real_part.end());

  // 计算虚部的平均值、方差、最大值和最小值
  this->sftf_imag_mean = std::accumulate(imag_part.begin(), imag_part.end(), 0.0f) / imag_part.size();
  float sum_of_squares_imag = std::inner_product(imag_part.begin(), imag_part.end(), imag_part.begin(), 0.0f);
  this->sftf_imag_variance = sum_of_squares_imag / imag_part.size();
  this->sftf_imag_min = *std::min_element(imag_part.begin(), imag_part.end());
  this->sftf_imag_max = *std::max_element(imag_part.begin(), imag_part.end());

  this->sftfRealArray = std::vector<float>(SFTF_NUM, 0);
  this->sftfImagArray = std::vector<float>(SFTF_NUM, 0);
  int index = 0;
  for (int i = 0; i < SFTF_NUM; i++) {
    while (real_part[index] == 0 && imag_part[index] == 0) {
      index++;
      if (index >= real_part.size() || index >= imag_part.size()) {
        index--;
        break;
      }
    }
    if (index < real_part.size() && index < imag_part.size()) {
      this->sftfRealArray[i] = real_part[index];
      this->sftfImagArray[i] = imag_part[index];
    }
  }
  Free(x);
}

void wavAudio::calculateMelspectrogram()
{
  // std::vector<float> x(this->samples_num);
  // std::vector<float> x(this->process_data.size());
  float* x = Malloc(float, this->samples_num);
  // std::transform(this->process_data.begin(), this->process_data.end(), x.begin(),
  std::transform(this->process_data, this->process_data+this->samples_num, x,
    [](int16_t a) {
      return static_cast<float>(a) / 32767.f;
    });

  // 计算梅尔频谱图
  int n_fft = 400;
  int n_hop = 160;
  int n_mel = 40;
  int fmin = 80;
  int fmax = 7600;

  //std::vector<std::vector<float>> mels = librosa::Feature::melspectrogram(x, this->sample_rate, n_fft, n_hop, "hann", true, "reflect", 2.f, n_mel, fmin, fmax);

  std::vector<float> x1;
  // for (int i = 0, index = 0; i < DATA_NUM && index < x.size(); index++) {
    for (int i = 0, index = 0; i < DATA_NUM && index < this->samples_num; index++) {
    if (x[index] != 0) {
      x1.push_back(x[index]);
      i++;
    }
  }
  std::vector<std::vector<float>> mels = melSpectrogram(x1, sample_rate, n_fft, n_hop, "hann", true, "reflect", 2.f, n_mel, fmin, fmax);

  // 将二维数组打平为一维数组
  std::vector<float> flattenedMels;
  for (const auto& row : mels) {
    flattenedMels.insert(flattenedMels.end(), row.begin(), row.end());
  }

  // 计算平均值
  this->mels_mean = std::accumulate(flattenedMels.begin(), flattenedMels.end(), 0.0f) / flattenedMels.size();

  // 计算方差
  float sum_of_squares = 0.0f;
  for (float value : flattenedMels) {
    sum_of_squares += (value - this->mels_mean) * (value - this->mels_mean);
  }
  this->mels_variance = sum_of_squares / flattenedMels.size();
  // 计算最小值和最大值
  this->mels_min = *std::min_element(flattenedMels.begin(), flattenedMels.end());
  this->mels_max = *std::max_element(flattenedMels.begin(), flattenedMels.end());

  // calculate the top MELS_NUM dimension
  this->melsArray = std::vector<float>(MELS_NUM, 0);
  int index = 0;
  for (int i = 0; i < MELS_NUM; i++) {
    while (flattenedMels[index] == 0)
    {
      index++;
      if (index >= flattenedMels.size()) {
        index--;
        break;
      }
    }
    if (index < flattenedMels.size()) {
      this->melsArray[i] = flattenedMels[index];
    }
  }
  Free(x);
}


bool wavAudio::computeCharacter()
{
  Serial.println("begin to calculate character...");
  bool res = this->processdata();
  Serial.println("this->processdata() after");
  if(!res){
    return false;
  }
  this->calculatePitch();
  Serial.println("this->calculatePitch() after");
  this->calculateAverageLoudness();
  Serial.println("this->calculateAverageLoudness() after");
  this->calculateTime();
  Serial.println("this->calculateTime() after");
  this->calculateEnergy();
  Serial.println("this->calculateEnergy() after");
  this->calculateSTFT();
  Serial.println("this->calculateSTFT() after");
  this->calculateMelspectrogram();
  Serial.println("calculate character done");
  return true;
}
  // 估算可用内存的函数
int remMemory() {
  char *buf;
  int size = 500000; // 可根据需要调整
  while ((buf = (char *)malloc(size -= 1000)) == NULL);
  free(buf);
  return size;
}

std::vector<float> wavAudio::returnCharacter()
{
  std::vector<float> res;

  Serial.print("res after Free memory: ");Serial.println(remMemory());  

  res.push_back(this->duration);
  res.push_back(this->pitch);
  res.push_back(this->average_loudness);
  res.push_back(this->energy);
  res.push_back(this->mels_mean);
  res.push_back(this->mels_variance);
  res.push_back(this->mels_min);
  res.push_back(this->mels_max);
  res.push_back(this->sftf_real_mean);
  res.push_back(this->sftf_real_variance);
  res.push_back(this->sftf_real_max);
  res.push_back(this->sftf_real_min);
  res.push_back(this->sftf_imag_mean);
  res.push_back(this->sftf_imag_variance);
  res.push_back(this->sftf_imag_max);
  res.push_back(this->sftf_imag_min);
  for (int i = 0; i < MELS_NUM; i++) {
    res.push_back(this->melsArray[i]);
  }
  for (int i = 0; i < SFTF_NUM; i++) {
    res.push_back(this->sftfRealArray[i]);
    // res.push_back(this->sftfImagArray[i]);
  }

  return res;
}

wavAudio::~wavAudio()
{
  Free(this->data);
  Free(this->process_data);
}
